import os
import shutil

import requests
import m3u8
from Crypto.Cipher import AES
from concurrent.futures import ThreadPoolExecutor
import subprocess
import json

"""
批量下载
读取fileinfo.json文件，key是视频文件名，value是视频的m3u8地址
开多线程下载视频，每个线程下载一个视频
每个视频的ts分块在不同的目录防止合并错乱（分块保存目录：ts_segments/xxx）
合并视频只合并自己的ts分块目录的
下载完成后，合并视频并删除分开，保存到指定目录（保存到当前同级的目录：videomp4），默认后缀名(.mp4)
"""

# 配置
FILEINFO_PATH = "2025下直播6.json"
TS_SEGMENTS_DIR = "ts_segments"
VIDEO_MP4_DIR = "videomp4"
THREADS = 10  # 下载线程数

# M3U8_URL = "https://example.com/path/to/video.m3u8"  # 替换为实际 M3U8 地址
SAVE_DIR = "./downloaded_video"
# THREADS = 10  # 下载线程数

# 创建保存目录
os.makedirs(TS_SEGMENTS_DIR, exist_ok=True)
os.makedirs(VIDEO_MP4_DIR, exist_ok=True)

# 1. 解析 M3U8 文件
def parse_m3u8(m3u8_url):
    playlist = m3u8.load(m3u8_url)
    key = None
    if playlist.keys and playlist.keys[0]:
        key_url = playlist.keys[0].uri
        base_url = m3u8_url.rsplit("/", 1)[0]
        key_url = base_url + "/" + key_url if not key_url.startswith("http") else key_url
        key_response = requests.get(key_url)
        key = key_url, key_response.content  # (key_url, key_bytes)
    ts_list = [seg.absolute_uri for seg in playlist.segments]
    return key, ts_list

# 2. 下载并解密 TS 分片
def download_ts(ts_url, index, key, ts_dir, iv=None):
    try:
        headers = {"User-Agent": "Mozilla/5.0"}
        response = requests.get(ts_url, headers=headers, stream=True)
        if response.status_code != 200:
            print(f"下载失败 {ts_url}")
            return

        data = response.content
        if key:
            cipher = AES.new(key[1], AES.MODE_CBC, iv=key[1] if iv is None else iv)
            data = cipher.decrypt(data)
        filename = os.path.join(ts_dir, f"segment_{index:04d}.ts")
        with open(filename, "wb") as f:
            f.write(data)
        print(f"下载完成 {filename}")
    except Exception as e:
        print(f"下载错误 {ts_url}: {e}")

# 3. 合并 TS 分片
def merge_ts_files(output_mp4, ts_dir):
    ts_files = sorted([os.path.join(ts_dir, f) for f in os.listdir(ts_dir) if f.endswith(".ts")])
    with open(os.path.join(ts_dir, "filelist.txt"), "w") as f:
        for ts in ts_files:
            f.write(f"file '{os.path.basename(ts)}'\n")

    # 使用 ffmpeg 合并
    command = [
        "ffmpeg",
        "-f", "concat",
        "-safe", "0",
        "-i", os.path.join(ts_dir, "filelist.txt"),
        "-c", "copy",
        output_mp4
    ]
    subprocess.run(command, check=True)
    print(f"合并完成: {output_mp4}")

# 主流程
def main():
    # 读取 fileinfo.json
    with open(FILEINFO_PATH, 'r', encoding='utf-8') as f:
        fileinfo = json.load(f)

    with ThreadPoolExecutor(max_workers=THREADS) as executor:
        for item in fileinfo:
            for video_name, m3u8_url in item.items():
                executor.submit(download_and_merge_video, video_name, m3u8_url)


def download_and_merge_video(video_name, m3u8_url):
    try:
        ts_dir = os.path.join(TS_SEGMENTS_DIR, video_name)
        os.makedirs(ts_dir, exist_ok=True)

        key, ts_list = parse_m3u8(m3u8_url)
        print(f"{video_name} 共找到 {len(ts_list)} 个 TS 分片")

        # 多线程下载
        with ThreadPoolExecutor(max_workers=THREADS) as executor:
            for i, ts_url in enumerate(ts_list):
                executor.submit(download_ts, ts_url, i, key, ts_dir)

        # 合并
        output_mp4 = os.path.join(VIDEO_MP4_DIR, f'{video_name}.mp4')
        merge_ts_files(output_mp4, ts_dir)

        # 合并并清理
        shutil.rmtree(ts_dir)
        print("TS 分片目录已删除")

        print(f"{video_name} 下载合并完成")
    except Exception as e:
        print(f"{video_name} 下载合并失败: {e}")

if __name__ == "__main__":
    main()
